import unittest
import numpy
from python.base.init import FLAGS
from python.database.base.data_util import DataUtil


class DataUtilUnittest(unittest.TestCase):
  def setUp(self):
    FLAGS.holiday_json_path = "presto/data/holidays.json"
    FLAGS.corp_action_path = "presto/data/corp_actions.json"

  def test_holidays(self):
    tse_holidays = [20110103, 20121231, 20161231, 20171231]
    for tse_holiday in tse_holidays:
      self.assertTrue(tse_holiday in DataUtil.get_holidays("TSE"))

    krx_holidays = [20111003, 20120101, 20120411]
    for krx_holiday in krx_holidays:
      self.assertTrue(krx_holiday in DataUtil.get_holidays("KRX"))

    # Korean SAT, it's not day off.
    self.assertFalse(20111110 in DataUtil.get_holidays("KRX"))

  def test_next_bd(self):
    self.assertEqual(DataUtil.get_next_bd(20160720, "KRX"), 20160721)
    self.assertEqual(DataUtil.get_next_bd(20160721, "KRX"), 20160722)
    self.assertEqual(DataUtil.get_days_next_bd(20160720, "KRX", 3), 20160725)
    self.assertEqual(DataUtil.get_next_bd(20160720, "TSE"), 20160721)
    self.assertEqual(DataUtil.get_next_bd(20160721, "TSE"), 20160722)
    self.assertEqual(DataUtil.get_days_next_bd(20160720, "TSE", 3), 20160725)

    self.assertEqual(DataUtil.get_next_bd(20110102, "TSE"), 20110104)
    self.assertEqual(DataUtil.get_next_bd(20160719, "TSE", False), 20160715)
    self.assertEqual(DataUtil.get_next_bd(20160719, "KRX", False), 20160718)

  def test_corp_action_ban_matrix(self):
    # Real data set.
    dates = DataUtil.get_range_bds(20160711, 20160930, "TSE")

    symbols = ["TSE-7262", "TSE-8270", "TSE-6191"]
    ban_matrix = DataUtil.get_corp_action_ban_matrix("TSE", dates, symbols, corp_actions=None)
    self.assertEqual(ban_matrix.dtype, numpy.bool)

    self.assertTrue(
        all(ban_matrix[dates.index(20160721):dates.index(20160727), symbols.index("TSE-7262")]))

    self.assertFalse(ban_matrix[dates.index(20160720), symbols.index("TSE-7262")])

    self.assertTrue(
        all(ban_matrix[dates.index(20160823):dates.index(20160829), symbols.index("TSE-8270")]))

    self.assertFalse(ban_matrix[dates.index(20160815), symbols.index("TSE-8270")])

    # Fake data set
    dates = DataUtil.get_range_bds(20160711, 20160930, "TSE")

    symbols = ["a", "b", "c", "d", "f"]
    artificial_corp_acts = [{
        "stop_date": "20160729", "symbol": "a", "event_type": "PAUSE", "resume_date": "20160802"
    },
                            {
                                "stop_date": "20160804",
                                "symbol": "a",
                                "event_type": "PAUSE",
                                "resume_date": "20160805"
                            },
                            {
                                "stop_date": "20160803",
                                "symbol": "b",
                                "event_type": "PAUSE",
                                "resume_date": "20160811"
                            },
                            {
                                "stop_date": "20161010",
                                "symbol": "c",
                                "event_type": "PAUSE",
                                "resume_date": "20161014"
                            },
                            {
                                "stop_date": "20160926",
                                "symbol": "d",
                                "event_type": "PAUSE",
                                "resume_date": "20160930"
                            },
                            {
                                "stop_date": "20160830",
                                "symbol": "f",
                                "event_type": "PAUSE",
                                "resume_date": "20160929"
                            }]

    ban_matrix = DataUtil.get_corp_action_ban_matrix("TSE", dates, symbols, artificial_corp_acts)

    self.assertEqual(ban_matrix.dtype, numpy.bool)

    self.assertTrue(all(ban_matrix[dates.index(20160729):dates.index(20160802),
                                   symbols.index("a")]))

    self.assertTrue(all(ban_matrix[dates.index(20160804):dates.index(20160805),
                                   symbols.index("a")]))

    self.assertFalse(ban_matrix[dates.index(20160727), symbols.index("a")])

    self.assertFalse(ban_matrix[dates.index(20160801), symbols.index("b")])

    self.assertFalse(ban_matrix[-1, symbols.index("c")])

    self.assertTrue(ban_matrix[-1, symbols.index("d")])

    self.assertTrue(all(ban_matrix[dates.index(20160830):dates.index(20160929),
                                   symbols.index("f")]))

    self.assertFalse(ban_matrix[-1, symbols.index("f")])


if __name__ == "__main__":
  unittest.main()
